import matplotlib.pyplot as plt
import seaborn as sns
import torch
import numpy as np
from sklearn.manifold import TSNE
from sklearn.metrics import roc_curve, auc
import os

# -------------------- 修正字体配置：删除错误的数学字体设置，简化配置 --------------------
plt.rcParams.update({
    'text.usetex': False,  # 禁用 LaTeX，彻底避免衬线字体依赖
    'axes.unicode_minus': False,  # 解决负号显示问题
    # 仅指定核心字体列表，覆盖 sans-serif 和 serif（用系统已安装字体）
    'font.sans-serif': ['SimSun', 'Microsoft YaHei', 'Arial', 'DejaVu Sans'],
    'font.serif': ['SimSun', 'Microsoft YaHei', 'Arial', 'DejaVu Sans'],
})

# -----------------------------------------------------------------------------

def plot_reconstruction_comparison(real, fake, save_path, n=5):
    real = real[:n].cpu().numpy().transpose(0, 2, 3, 1)
    fake = fake[:n].cpu().numpy().transpose(0, 2, 3, 1)
    plt.figure(figsize=(15, 6))
    for i in range(n):
        plt.subplot(2, n, i+1)
        plt.imshow((real[i] * 0.5 + 0.5).clip(0, 1))
        plt.title("Real", fontsize=10)
        plt.axis("off")
        plt.subplot(2, n, i+1 + n)
        plt.imshow((fake[i] * 0.5 + 0.5).clip(0, 1))
        plt.title("Reconstructed", fontsize=10)
        plt.axis("off")
    plt.tight_layout()
    plt.savefig(save_path, dpi=100, bbox_inches='tight')
    plt.close()

def plot_latent_distribution(latent_i, latent_o, labels, save_path):
    latent_i = latent_i.cpu().numpy()
    latent_o = latent_o.cpu().numpy()
    labels = labels.cpu().numpy()
    all_latent = np.vstack([latent_i, latent_o])
    all_labels = np.hstack([labels, labels])
    all_types = np.hstack([np.zeros(len(latent_i)), np.ones(len(latent_o))])

    tsne = TSNE(n_components=2, perplexity=30)
    latent_2d = tsne.fit_transform(all_latent)

    plt.figure(figsize=(10, 8))
    sns.scatterplot(
        x=latent_2d[:, 0], y=latent_2d[:, 1],
        hue=all_labels, style=all_types,
        palette={0: "green", 1: "red"},
        markers={0: "o", 1: "s"},
        alpha=0.7
    )
    plt.title("Latent Space Distribution (TSNE)", fontsize=12)
    plt.legend(
        title="Label (0=Normal, 1=Anomaly)\nType (0=Input, 1=Output)",
        fontsize=9
    )
    plt.tight_layout()
    plt.savefig(save_path, dpi=100, bbox_inches='tight')
    plt.close()

def plot_anomaly_score_distribution(scores, labels, save_path):
    scores = scores.cpu().numpy()
    labels = labels.cpu().numpy()
    plt.figure(figsize=(8, 5))
    sns.histplot(
        data={"score": scores, "label": labels},
        x="score", hue="label", multiple="stack",
        bins=30, palette={0: "green", 1: "red"}
    )
    normal_scores = scores[labels == 0]
    threshold = np.mean(normal_scores) + 3 * np.std(normal_scores) if len(normal_scores) > 0 else 0.5
    plt.axvline(x=threshold, color="black", linestyle="--", label="Threshold")
    plt.title("Anomaly Score Distribution", fontsize=12)
    plt.xlabel("Anomaly Score (L2 Distance of Latent Vectors)", fontsize=10)
    plt.ylabel("Count", fontsize=10)
    plt.legend(title="Label", labels=["Normal", "Anomaly", "Threshold"], fontsize=9)
    plt.tight_layout()
    plt.savefig(save_path, dpi=100, bbox_inches='tight')
    plt.close()




def plot_roc_curve(labels, scores, save_path):
    """绘制ROC曲线"""
    fpr, tpr, _ = roc_curve(labels.cpu().numpy(), scores.cpu().numpy())
    roc_auc = auc(fpr, tpr)

    plt.figure(figsize=(8, 6))
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {roc_auc:.3f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.grid(True, alpha=0.3)
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()


def plot_score_distribution_by_component(scores_dict, labels, save_path):
    """绘制各分量的分数分布"""
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    axes = axes.flatten()

    normal_mask = labels == 0
    anomaly_mask = labels == 1

    for idx, (name, scores) in enumerate(scores_dict.items()):
        if idx >= 4:
            break

        ax = axes[idx]
        if len(scores[normal_mask]) > 0:
            ax.hist(scores[normal_mask].cpu().numpy(), bins=50, alpha=0.5, label='Normal', density=True)
        if len(scores[anomaly_mask]) > 0:
            ax.hist(scores[anomaly_mask].cpu().numpy(), bins=50, alpha=0.5, label='Anomaly', density=True)

        ax.set_title(f'{name} Score Distribution')
        ax.set_xlabel('Score')
        ax.set_ylabel('Density')
        ax.legend()
        ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()